/*
 * @lc app=leetcode.cn id=337 lang=cpp
 *
 * [337] 打家劫舍 III
 */
#include "include.h"

struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode() : val(0), left(nullptr), right(nullptr) {}
    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
    TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
};

// @lc code=start
/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
public:
    int rob(TreeNode* root) {
        // todo : 用深度优先搜索的办法后序遍历这棵二叉树 ~
        // first right, then left
        // 甚至不需要DP空间

        // todo : better =>
        // DP[current_root][not_robbed] = 
        // max(DP[left_son][is_robbed], DP[left_son][not_robbed]) +
        // max(DP[right_son][is_robbed], DP[right_son][not_robbed])

        if (root == nullptr){return 0;}
        // DP ?
        // DP[current_root][is_robbed] = DP[left_son][not_robbed] + DP[right_son][not_robbed] + current_root->val
        // DP[current_root][not_robbed] = max(
        //     DP[left_son][is_robbed] + DP[right_son][is_robbed],
        //     DP[left_son][not_robbed] + DP[right_son][is_robbed],
        //     DP[left_son][is_robbed] + DP[right_son][not_robbed],
        //     DP[left_son][not_robbed] + DP[right_son][not_robbed]
        // )
        //
        // DP[current_root][finally] = max(DP[current_root][is_robbed], DP[current_root][not_robbed])

        levelTraverse(root);

        reverseDP();

        return std::max(DP_[0].is_robbed, DP_[0].not_robbed);
    }
private:    
    struct NodeInfo{
        int selfIndex = -1;

        int leftSonIndex = -1;
        int rightSonIndex = -1;

        int is_robbed = 0;
        int not_robbed = 0;
        int self_val = 0;
    };
    std::vector<NodeInfo> DP_;
private:
    inline void levelTraverse(TreeNode* root){
        std::queue<TreeNode*> levelQueue;
        levelQueue.push(root);

        int nodeIndex = 0;
        DP_.emplace_back(NodeInfo());

        TreeNode* currentNode;
        while(!levelQueue.empty()){
            currentNode = levelQueue.front();
            levelQueue.pop();

            DP_[nodeIndex].selfIndex = nodeIndex;
            DP_[nodeIndex].self_val = currentNode->val;

            if (currentNode->left){
                levelQueue.push(currentNode->left);
            
                DP_.emplace_back(NodeInfo());
                DP_[nodeIndex].leftSonIndex = DP_.size() - 1;
            }
            if (currentNode->right){
                levelQueue.push(currentNode->right);

                DP_.emplace_back(NodeInfo());
                DP_[nodeIndex].rightSonIndex = DP_.size() - 1;
            }

            nodeIndex += 1;
        }
    }

    inline void reverseDP(){
        for (int i=DP_.size()-1;i>=0;i--){
            auto& info = DP_[i];
            // leaf
            if (info.leftSonIndex == -1 and info.rightSonIndex == -1){
                info.is_robbed = info.self_val;
                // info.not_robbed = 0;
            }
            // not leaf, has only left son
            if (info.leftSonIndex != -1 and info.rightSonIndex == -1){
                info.is_robbed = DP_[info.leftSonIndex].not_robbed + info.self_val;
                info.not_robbed = std::max(
                    DP_[info.leftSonIndex].is_robbed,
                    DP_[info.leftSonIndex].not_robbed
                );
            }
            // not leaf, has only right son
            if (info.leftSonIndex == -1 and info.rightSonIndex != -1){
                info.is_robbed = DP_[info.rightSonIndex].not_robbed + info.self_val;
                info.not_robbed = std::max(
                    DP_[info.rightSonIndex].is_robbed,
                    DP_[info.rightSonIndex].not_robbed
                );
            }
            // not leaf, has two son
            if (info.leftSonIndex != -1 and info.rightSonIndex != -1){
                info.is_robbed = DP_[info.leftSonIndex].not_robbed + DP_[info.rightSonIndex].not_robbed + info.self_val;
                info.not_robbed = std::max(
                    std::max(
                        DP_[info.leftSonIndex].is_robbed + DP_[info.rightSonIndex].is_robbed,
                        DP_[info.leftSonIndex].not_robbed + DP_[info.rightSonIndex].is_robbed
                    ),
                    std::max(
                        DP_[info.leftSonIndex].is_robbed + DP_[info.rightSonIndex].not_robbed,
                        DP_[info.leftSonIndex].not_robbed + DP_[info.rightSonIndex].not_robbed
                    )
                );
            }
        }
    }
};
// @lc code=end

